# -*- coding: utf-8 -*-
"""
Created on Wed Oct 27 19:06:38 2021

@author: singularity

参考博客：
https://blog.csdn.net/qq_42109740/article/details/105780340

YOLO损失函数形式：
loss = l * 矩形框中心位置交叉熵 （有效目标）
        + l * 框的宽高误差MSE（有效目标）
         + 置信度交叉熵（有效目标，无效目标）
         + 类别交叉熵（有效目标）

来自网络输出预测量形式：
(Tensor[B,3*(1+4+5),16,16],Tensor[B,3*(1+4+5),32,32])
  ↑大目标特征图              ↑小目标特征图

Box内容：1置信度 + 4中心XY偏移、长宽WH系数 + 5类别

来自ground truth数据形式：
(Tensor[B,3*(1+4+5 +1),16,16], Tensor[B,3*(1+4+5 +1),32,32])
                   ↑目标归属逻辑值            ↑置信度、位置信息与类别概率
"""

import torch
import torchvision
import torchvision.transforms as transforms

import torch.nn as nn
import torch.nn.functional as F

class YOLOJ_Loss(nn.Module):
    def __init__(self, classN:int, AnchorN:int, lambdNoobj:float=0.5):
        super(YOLOJ_Loss,self).__init__()
        
        self.classN = classN    # 分辨的类别
        self.AnchrN = AnchorN   # 每个特征图通道的Anchor数目
        
        self.lbdNoobj = lambdNoobj  # 无效对象的权值
        
        self.BCE_X = nn.BCELoss(reduction='none')
        self.BCE_Y = nn.BCELoss(reduction='none')
        self.BCE_P = nn.BCELoss(reduction='none')
        self.BCE_C = nn.BCELoss(reduction='none')
    
    def forward(self, pred:(torch.Tensor,torch.Tensor), 
                    gt:(torch.Tensor,torch.Tensor),
                    debug:bool = False):
        if debug:
            print(f"输入数据shape：({pred[0].shape},{pred[1].shape}) 与 ({gt[0].shape},{gt[1].shape})")
        
        lossing = 0
        lossingNow = 0
        
        for layer in range(2):
            if debug:
                print(f"特征图{layer}")
            for anchorId in range(self.AnchrN):
                if debug:
                    print(f"--Anchor {anchorId}--")
                    
                ''' 先求有效目标大小调整系数'''
                lbd_coord = gt[layer][:,anchorId * 11 +3] * \
                                gt[layer][:,anchorId * 11 +4]   # 真实长宽
                lbd_coord = 2- lbd_coord
                
                if debug:
                    print(f"目标大小平均系数shape：{lbd_coord.shape}")
                
                ''' 计算中心偏移的交叉熵 '''
                lossingNow = self.BCE_X(pred[layer][:,anchorId*10 + 1],
                                        gt[layer][:,anchorId*11 + 1])
                lossingNow += self.BCE_Y(pred[layer][:,anchorId*10 + 2],
                                         gt[layer][:,anchorId*11 + 2])

                # -- lambda*Iobj*Entrophy
                lossingNow = lbd_coord * \
                                gt[layer][:,anchorId * 11 +10] * \
                                    lossingNow
                
                if debug:
                    print(f"中心偏移交叉熵shape：{lossingNow.shape}")
                
                lossing = torch.sum(torch.sum(lossingNow , dim = 1),dim = 1)
                    
                ''' 计算长宽误差MSE'''
                lossingNow = torch.pow( gt[layer][:,anchorId*11 +3]-pred[layer][:,anchorId*10 + 3] ,
                                       2)
                lossingNow += torch.pow( gt[layer][:,anchorId*11 +4]-pred[layer][:,anchorId*10 + 4] ,
                                       2)
                
                if debug:
                    print(f"长宽MSE shape：{lossingNow.shape}")
                
                # --lambda*Iobj*SE
                lossingNow = torch.sum( lbd_coord * gt[layer][:,anchorId * 11 +10] * lossingNow,
                                       dim = 1)
                
                lossing += 0.5 * torch.sum(lossingNow,dim = 1)
                
                ''' 置信度交叉熵'''
                lossingNow = self.BCE_P(pred[layer][:,anchorId*10],
                                        gt[layer][:,anchorId*11])
                
                if debug:
                    print(f"置信度交叉熵shape：{lossingNow.shape}")
                    
                    
                # -- 有效目标
                lossing += torch.sum(torch.sum( gt[layer][:,anchorId * 11 +10] * lossingNow,
                                               dim = 1),
                                     dim =1 )
                
                # -- 无效目标
                lossing += self.lbdNoobj * torch.sum(torch.sum( (1-gt[layer][:,anchorId * 11 +10]) * lossingNow,
                                                               dim = 1),
                                                     dim =1 )
                
                ''' 类别交叉熵'''                
                lossingNow = self.BCE_C(pred[layer][:,anchorId*10+5:anchorId*10+10],
                                        gt[layer][:,anchorId*11+5:anchorId*11+10])
                
                if debug:
                    print(f"类别交叉熵shape：{lossingNow.shape}")
                
                lossingNow = torch.sum(lossingNow,dim = 1)  # 类别积分
                
                lossing += torch.sum(torch.sum( gt[layer][:,anchorId * 11 +10] * lossingNow,
                                               dim = 1),
                                     dim =1 )
                
        # 计算batch的平均loss
        
        if debug:
            print(f"Batch loss shape：{lossing.shape}")
                
        lossing = torch.mean(lossing) /2.0
        return lossing
                
if __name__ == '__main__':
    BatchSize = 8
    
    # 实例化损失
    criterion = YOLOJ_Loss(5, 3, 0.5)

    # 创建随机数据
    predT = ( torch.rand((BatchSize,30,16,16)), torch.rand((BatchSize,30,32,32)) )
    gT = ( torch.rand((BatchSize,33,16,16)), torch.rand((BatchSize,33,32,32)) )
    
    # 计算损失
    loss = criterion(predT,gT)